# 导入必要库
import warnings
from collections import Counter  # 用于模型融合的投票统计

import joblib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from catboost import CatBoostClassifier  # 集成学习模型
from flaml import AutoML  # 自动化机器学习工具
from lightgbm import LGBMClassifier  # 轻量级梯度提升机
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.decomposition import PCA
from sklearn.metrics import accuracy_score, classification_report  # 评估指标
from sklearn.model_selection import KFold  # 交叉验证工具
from sklearn.pipeline import Pipeline  # 管道工具（当前未直接使用）
from sklearn.preprocessing import LabelEncoder, StandardScaler  # 标签编码和标准化
from xgboost import XGBClassifier  # 极端梯度提升机

warnings.filterwarnings('ignore')  # 关闭警告提示

# 配置flaml日志级别（避免训练信息过多）
import logging

# 设置flaml的日志级别为WARNING（仅显示警告及以上日志）
logging.getLogger('flaml.automl.logger').setLevel(logging.WARNING)

# 加载数据（训练集、测试集、提交样例）
train = pd.read_csv('./train.csv', index_col='id')  # 训练集（含目标变量）
test = pd.read_csv('./test.csv', index_col='id')    # 测试集（需预测）
sub = pd.read_csv('./sample_submission.csv', index_col='id')  # 提交格式样例


# 第一步：数据概览（查看数据基本信息）
def load_and_overview_data(train, test):
    """打印数据基本信息（形状、列名、缺失值、目标分布）"""
    print('Train shape:', train.shape)  # 训练集维度（样本数×特征数）
    print('Test shape:', test.shape)    # 测试集维度
    print('\nTrain columns:', train.columns.tolist())  # 训练集特征列表
    print('\nTrain info:')
    print(train.info())  # 数据类型和非空值统计
    print('\nTrain missing values:')
    print(train.isnull().sum())  # 各特征缺失值数量
    print('\nTrain target value counts:')
    print(train['Personality'].value_counts())  # 目标变量类别分布


# 第二步：数据预处理（处理缺失值和类别特征）
def preprocess_data(train, test):
    """
    处理数值型和类别型特征的缺失值，统一类别格式
    :return: 处理后的训练集、测试集、填充信息（用于复现）
    """
    num_cols = ['Time_spent_Alone', 'Social_event_attendance', 'Going_outside', 'Friends_circle_size', 'Post_frequency']  # 数值特征列表
    cat_cols_train = ['Stage_fear', 'Drained_after_socializing', 'Personality']  # 训练集类别特征（含目标）
    cat_cols_test = ['Stage_fear', 'Drained_after_socializing']  # 测试集类别特征（不含目标）
    
    # 数值特征用训练集中位数填充（避免数据泄露）
    num_medians = train[num_cols].median()
    train[num_cols] = train[num_cols].fillna(num_medians)
    test[num_cols] = test[num_cols].fillna(num_medians)
    
    # 类别特征统一格式（转小写、去空格）并填充缺失值为'missing'
    for col in cat_cols_train:
        train[col] = train[col].astype(str).str.strip().str.lower()  # 统一格式
        train[col] = train[col].replace('nan', np.nan).fillna('missing')  # 填充缺失值
    for col in cat_cols_test:
        test[col] = test[col].astype(str).str.strip().str.lower()
        test[col] = test[col].replace('nan', np.nan).fillna('missing')
    
    # 记录填充信息（用于后续可能的验证）
    fill_info = {'num_medians': num_medians.to_dict(), 'cat_fill': 'missing'}
    return train, test, fill_info


# --- 增强版特征工程 ---
def build_features(train, test, target_col='Personality',
                   add_interactions=True, add_cat_combo=True, add_bins=True,
                   add_stats=True, select_top_n=20, use_pca=True, pca_n=8,
                   use_target_encoding=True, use_freq_encoding=True):
    """
    构造模型输入特征（含多种增强特征工程）
    :return: 训练集特征、训练集目标、测试集特征
    """
    """构造模型输入特征（含多种增强特征工程）"""
    num_cols = ['Time_spent_Alone', 'Social_event_attendance', 'Going_outside', 'Friends_circle_size', 'Post_frequency']
    cat_cols = ['Stage_fear', 'Drained_after_socializing']
    df_train = train.copy()
    df_test = test.copy()

    # 1. 数值特征交互
    if add_interactions:
        for i in range(len(num_cols)):
            for j in range(i+1, len(num_cols)):
                col_name = f'{num_cols[i]}_x_{num_cols[j]}'
                df_train[col_name] = df_train[num_cols[i]] * df_train[num_cols[j]]
                df_test[col_name] = df_test[num_cols[i]] * df_test[num_cols[j]]

    # 2. 类别特征组合
    if add_cat_combo:
        combo_col = cat_cols[0] + '_' + cat_cols[1]
        df_train[combo_col] = df_train[cat_cols[0]].astype(str) + '_' + df_train[cat_cols[1]].astype(str)
        df_test[combo_col] = df_test[cat_cols[0]].astype(str) + '_' + df_test[cat_cols[1]].astype(str)
        cat_combo_cols = cat_cols + [combo_col]
    else:
        cat_combo_cols = cat_cols.copy()

    # 3. 连续特征分箱（等频分箱）
    if add_bins:
        for col in num_cols:
            # 基于训练集确定分箱边界（等频）
            _, bins = pd.qcut(df_train[col], 4, duplicates='drop', retbins=True)
            # 测试集使用训练集的分箱边界
            df_train[f'{col}_bin'] = pd.cut(df_train[col], bins=bins, labels=False, include_lowest=True)
            df_test[f'{col}_bin'] = pd.cut(df_test[col], bins=bins, labels=False, include_lowest=True)
        bin_cols = [f'{col}_bin' for col in num_cols]
    else:
        bin_cols = []

    # 4. 统计特征
    if add_stats:
        # 全局相对统计（当前值 vs 全局统计量）
        for col in num_cols:
            global_mean = df_train[col].mean()
            global_max = df_train[col].max()
            global_min = df_train[col].min()
            # 新增相对特征（当前值与全局统计量的差异）
            df_train[f'{col}_vs_mean'] = df_train[col] - global_mean
            df_train[f'{col}_vs_max'] = df_train[col] - global_max
            df_train[f'{col}_vs_min'] = df_train[col] - global_min
            # 测试集使用训练集的全局统计量
            df_test[f'{col}_vs_mean'] = df_test[col] - global_mean
            df_test[f'{col}_vs_max'] = df_test[col] - global_max
            df_test[f'{col}_vs_min'] = df_test[col] - global_min

        # 分组统计（保留原逻辑）
        for cat in cat_cols:
            for col in num_cols:
                grp_mean = df_train.groupby(cat)[col].transform('mean')
                df_train[f'{col}_mean_by_{cat}'] = grp_mean
                mean_map = df_train.groupby(cat)[col].mean().to_dict()
                df_test[f'{col}_mean_by_{cat}'] = df_test[cat].map(mean_map)


    # 5. 类别特征编码优化
    # 5.1 目标编码（使用交叉验证折计算OOB编码）
    if use_target_encoding:
        kf_enc = KFold(n_splits=5, shuffle=True, random_state=42)  # 与后续验证折数一致
        for cat in cat_cols:
            df_train[f'{cat}_target_enc'] = 0.5  # 初始化默认值
            # 按折计算OOB编码
            for train_enc_idx, val_enc_idx in kf_enc.split(df_train):
                # 用当前训练折计算目标频率
                target_map = df_train.iloc[train_enc_idx].groupby(cat)[target_col].apply(
                    lambda x: x.value_counts(normalize=True).get('extrovert', 0.5)
                )
                # 填充当前验证折的编码值
                df_train.loc[val_enc_idx, f'{cat}_target_enc'] = df_train.iloc[val_enc_idx][cat].map(target_map).fillna(0.5)
            # 测试集用全量训练数据的目标编码（最终预测使用）
            target_map_full = df_train.groupby(cat)[target_col].apply(
                lambda x: x.value_counts(normalize=True).get('extrovert', 0.5)
            )
            df_test[f'{cat}_target_enc'] = df_test[cat].map(target_map_full).fillna(0.5)

    # 5.2 频率编码（同理改为OOB）
    if use_freq_encoding:
        kf_enc = KFold(n_splits=5, shuffle=True, random_state=42)
        for cat in cat_cols:
            df_train[f'{cat}_freq_enc'] = 0.0  # 初始化默认值
            # 按折计算OOB频率
            for train_enc_idx, val_enc_idx in kf_enc.split(df_train):
                freq_map = df_train.iloc[train_enc_idx][cat].value_counts(normalize=True)
                df_train.loc[val_enc_idx, f'{cat}_freq_enc'] = df_train.iloc[val_enc_idx][cat].map(freq_map).fillna(0)
            # 测试集用全量训练数据的频率
            freq_map_full = df_train[cat].value_counts(normalize=True)
            df_test[f'{cat}_freq_enc'] = df_test[cat].map(freq_map_full).fillna(0)

    # 6. 数值特征标准化（扩展到所有数值增强特征）
    # 原始数值特征 + 交互特征 + 统计特征
    all_num_cols = num_cols + [c for c in df_train.columns if any(s in c for s in ['_x_', '_vs_', '_by_'])]
    scaler = StandardScaler()
    train_num = scaler.fit_transform(df_train[all_num_cols])
    test_num = scaler.transform(df_test[all_num_cols])

    # 7. 类别特征独热编码（含组合、分箱）
    from sklearn.preprocessing import OneHotEncoder

    cat_cols_used = cat_combo_cols + bin_cols
    # 初始化OneHotEncoder（忽略测试集新类别）
    ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
    train_cat = ohe.fit_transform(df_train[cat_cols_used])  # 输出为NumPy数组
    test_cat = ohe.transform(df_test[cat_cols_used])        # 输出为NumPy数组


    # 8. 其余增强特征
    extra_cols = [c for c in df_train.columns if any(s in c for s in ['_x_', '_mean', '_max', '_min', '_by_', '_enc'])]
    train_extra = df_train[extra_cols].values if extra_cols else np.empty((len(df_train), 0))
    test_extra = df_test[extra_cols].values if extra_cols else np.empty((len(df_test), 0))

    # 9. 合并所有特征
    X_train = np.hstack([train_num, train_cat, train_extra])
    X_test = np.hstack([test_num, test_cat, test_extra])
    y_train = df_train[target_col].values  # 关键修正：提前赋值y_train


    # 10. 特征选择（LGBM重要性筛选）
    if select_top_n is not None and X_train.shape[1] > select_top_n:
        lgbm = LGBMClassifier(n_estimators=100, random_state=42)
        lgbm.fit(X_train, y_train)  # 此时y_train已定义
        importances = lgbm.feature_importances_
        idx = np.argsort(importances)[::-1][:select_top_n]
        X_train = X_train[:, idx]
        X_test = X_test[:, idx]

    # 11. PCA降维
    if use_pca and X_train.shape[1] > pca_n:
        pca = PCA(n_components=pca_n, random_state=42)
        X_train = pca.fit_transform(X_train)
        X_test = pca.transform(X_test)

    return X_train, y_train, X_test


# -------------------- 数据处理流程 --------------------
# 1. 数据概览（打印基本信息）
load_and_overview_data(train, test)

# 2. 预处理数据（填充缺失值）
train, test, fill_info = preprocess_data(train, test)

# 3. 特征工程（构造模型输入）
X_train, y_train, X_test = build_features(train, test, target_col='Personality')

# 4. 目标变量标签编码（将文本类别转为数值）
le = LabelEncoder()
y_train_enc = le.fit_transform(y_train)  # 训练集目标编码


# -------------------- AutoML 自动化训练 --------------------
print('\nAutoML 训练开始...')
# 提升AutoML计算复杂度和搜索空间
settings = {
    "time_budget": 360,  # 增加训练时间预算（秒）
    "task": 'classification',
    "log_file_name": 'flaml.log',
    "metric": 'accuracy',
    "estimator_list": [
        'lgbm', 'xgboost', 'catboost', 'rf', 'extra_tree', 'lrl1', 'lrl2'
    ],  # 加入更多模型
    "n_jobs": -1,  # 多核并行
    "seed": 42,
    "verbose": 1,
    # 可选：可进一步自定义搜索空间
    # "auto_augment": True,
}
automl = AutoML()
automl.fit(X_train=X_train, y_train=y_train_enc, **settings)  # 启动自动训练

# 输出AutoML结果
print('最佳模型：', automl.model)          # 自动选择的最优模型
print('最佳参数：', automl.best_config)    # 最优模型的超参数
print('最佳分数：', automl.best_loss)      # 最优模型的损失值

# 生成测试集预测（编码后转回原始类别）
automl_preds_enc = automl.predict(X_test)
automl_preds = le.inverse_transform(automl_preds_enc)  # 逆编码得到原始类别

# 训练集评估（验证过拟合情况）
automl_train_preds_enc = automl.predict(X_train)
automl_train_preds = le.inverse_transform(automl_train_preds_enc)
print('\n训练集评估：')
print(classification_report(y_train, automl_train_preds))  # 分类报告（精确率、召回率等）
print('训练集准确率：', accuracy_score(y_train, automl_train_preds))  # 准确率

# 保存AutoML预测结果（调整类别首字母大写，符合提交格式）
df_automl_pred = pd.DataFrame({'id': test.index, 'Personality': automl_preds})
df_automl_pred['Personality'] = df_automl_pred['Personality'].str.capitalize()
# df_automl_pred.to_csv('submission_automl.csv', index=False)


# -------------------- LGBM 模型（KFold交叉验证） --------------------
print('\nLGBM 模型训练开始...')
# 定义LGBM参数（手动调参）
lgbm_params = {
    'n_estimators': 200,    # 迭代次数
    'learning_rate': 0.05,  # 学习率
    'max_depth': 7,         # 树最大深度
    'random_state': 42      # 随机种子（保证可复现）
}
model = LGBMClassifier(**lgbm_params)

# 5折交叉验证（评估模型泛化能力）
kf = KFold(n_splits=5, shuffle=True, random_state=42)  # 5折，随机打乱数据
lgbm_oof_preds = np.zeros(len(X_train), dtype=int)     # 保存各折验证集预测结果
lgbm_oof_probs = np.zeros((len(X_train), len(le.classes_)))  # 保存概率值
fold_scores = []  # 保存各折准确率

for fold, (train_idx, val_idx) in enumerate(kf.split(X_train)):
    # 划分训练/验证集
    X_tr, X_val = X_train[train_idx], X_train[val_idx]
    y_tr, y_val = y_train_enc[train_idx], y_train_enc[val_idx]
    
    # 模型训练+验证
    model.fit(X_tr, y_tr)
    val_pred = model.predict(X_val)       # 验证集类别预测
    val_prob = model.predict_proba(X_val) # 验证集概率预测
    lgbm_oof_preds[val_idx] = val_pred    # 记录验证集预测结果
    lgbm_oof_probs[val_idx] = val_prob    # 记录概率值
    acc = accuracy_score(y_val, val_pred) # 计算当前折准确率
    print(f"Fold {fold+1} accuracy: {acc:.4f}")
    fold_scores.append(acc)

# 输出交叉验证结果
print("\n5折交叉验证平均准确率：", np.mean(fold_scores))
print("LGBM OOF分类报告：")
print(classification_report(y_train, le.inverse_transform(lgbm_oof_preds)))  # OOF预测与真实值对比

# 全量数据训练并生成测试集预测
model.fit(X_train, y_train_enc)
lgbm_preds_enc = model.predict(X_test)
lgbm_preds = le.inverse_transform(lgbm_preds_enc)  # 逆编码得到原始类别

# 保存LGBM预测结果
df_lgbm_pred = pd.DataFrame({'id': test.index, 'Personality': lgbm_preds})
df_lgbm_pred['Personality'] = df_lgbm_pred['Personality'].str.capitalize()
# df_lgbm_pred.to_csv('submission_lgbm.csv', index=False)


# -------------------- 二模型概率加权融合 --------------------
print('\n二模型概率加权融合开始...')
# 获取模型概率输出（需修改原代码保存概率）
# AutoML概率预测（假设AutoML支持predict_proba）
automl_probs = automl.predict_proba(X_test)
# LGBM概率预测（原代码已计算lgbm_oof_probs，全量训练后需重新计算）
model.fit(X_train, y_train_enc)
lgbm_probs = model.predict_proba(X_test)

# 保存AutoML最佳模型
joblib.dump(automl, 'automl_best_model.pkl')
# 保存LGBM最终模型
joblib.dump(model, 'lgbm_final_model.pkl')

# 加权融合（权重可通过验证集调优，示例用等权）
blend_probs = 0.5 * automl_probs + 0.5 * lgbm_probs
ensemble2_preds = le.inverse_transform(blend_probs.argmax(axis=1))  # 概率转类别

# 保存融合预测结果
df_ensemble2 = pd.DataFrame({'id': test.index, 'Personality': ensemble2_preds}) 
df_ensemble2['Personality'] = df_ensemble2['Personality'].str.capitalize()
df_ensemble2.to_csv('submission.csv', index=False)


# -------------------- 预测分布对比（分析模型差异） --------------------
print('\nAutoML测试集预测分布:')
print(pd.Series(automl_preds).value_counts())
print('\nLGBM+KFold测试集预测分布:')
print(pd.Series(lgbm_preds).value_counts())
print('\n二模型融合后测试集预测分布:')
print(pd.Series(ensemble2_preds).value_counts())


'''
主要修改了参数规模， 让 flaml搜索更多的空间。

'''